from nets import get_model_from_name
from OpCounter.thop import profile, clever_format
import torch


backbone = "convbk"
aa = ""

use_gpu = True if torch.cuda.is_available() else False

if backbone != "vit":
    model = get_model_from_name[backbone](num_classes=100, pretrained=False, aa=aa, use_gpu=use_gpu)
else:
    model = get_model_from_name[backbone](input_shape=[224, 224], num_classes=100, pretrained=False)

input = torch.randn(1, 3, 224, 224)
macs, params = profile(model, inputs=(input, ))
macs, params = clever_format([macs, params], "%.3f")

with open("evaluate_result/parameter.txt", "a") as f:
    f.write("{}{}=======>FLOPs:{}\tPara:{}\n".format(backbone, aa, macs, params))
    f.close()

print("FLOPs:{}\tPara:{}".format(macs, params))
